#include "Lapack.h"
#include <cmath>
#include <algorithm>
#include <float.h>
#include "../Exceptions.h"

using namespace std;

inline float sqr(float x) {
	return x * x;
}

void initIntervals(int n, Interval* intervals, Interval interval) {
	for (int i = 0; i < n; i++) {
		intervals[i].low = (i == 0) ? interval.low : interval.up;
		intervals[i].up = interval.up;
	}
}

int getNumEigensLess(const float *diagonal, 
		const float *offDiagonal, 
        int width, 
        float x)
{
    int count = 0;
    float prevDiff = (diagonal[0] - x);
    count += (prevDiff < 0)? 1 : 0;
    for (int i = 1; i < width; i++) {
        float diff = (diagonal[i] - x) - sqr(offDiagonal[i - 1]) / prevDiff;
		count += (diff < 0)? 1 : 0;
        prevDiff = diff;
    }
    return count;
}

void getNumEigensInInterval(int* numEigenIntervals,
                            Interval* intervals,
                            const float* diagonal, 
                            const float* offDiagonal,
                            int width)
{
	for (int i = 0; i < width; i++) {
		int numEigensLessThanLower = getNumEigensLess(diagonal, offDiagonal, width, intervals[i].low);
		int numEigensLessThanUpper = getNumEigensLess(diagonal, offDiagonal, width, intervals[i].up);

		numEigenIntervals[i] = numEigensLessThanUpper - numEigensLessThanLower;
	}
}

void recalcEigenIntervals(int n, 
		Interval* outIntervals,
		Interval* intervals,
		int  *numEigensInInterval,
		const float *d,
		const float *e,
		float tolerance)
{
	for (int threadId = 0; threadId < n; threadId++) {	
				
		int splitIndex = threadId;
		int inputIndex = 0;
		while (splitIndex >= numEigensInInterval[inputIndex]) {
			splitIndex -= numEigensInInterval[inputIndex];
			++inputIndex;
		}

		Interval src = intervals[inputIndex];
		float width = src.up - src.low;
		
		if (width <= tolerance) {
			// interval is narrow enough
			outIntervals[threadId] = src;			
		} else if (numEigensInInterval[inputIndex] == 1) {
			// if the number of eigenvalues in the interval is just 1
			// use interval bisection
			float middle = (src.up + src.low) / 2;
			int numEigens = getNumEigensLess(d, e, n, middle);
			numEigens -= getNumEigensLess(d, e, n, src.low);
						
			outIntervals[threadId].low = (numEigens != 0) ? src.low : middle;
			outIntervals[threadId].up = (numEigens != 0) ? middle : src.up;
		} else {
			// split the intervals into equal intervals of size divisionWidth
			float divisionWidth = width / numEigensInInterval[inputIndex];
			outIntervals[threadId].low = src.low + divisionWidth * splitIndex;
			outIntervals[threadId].up = outIntervals[threadId].low + divisionWidth;
		}
	}
}

Interval Lapack::getGerschgorinInterval(
	const float * diagonal, const float * offDiagonal, int length)
{
	Interval bound;
	bound.low = diagonal[0] - abs(offDiagonal[0]);
	bound.up = diagonal[0] + abs(offDiagonal[0]);    

    for (int i = 1; i < length - 1; i++) {
        float r =  abs(offDiagonal[i - 1]) + abs(offDiagonal[i]);
		bound.low = min(bound.low, diagonal[i] - r);
		bound.up = max(bound.up, diagonal[i] + r);
    }

	bound.low = min(bound.low, diagonal[length - 1] - abs(offDiagonal[length - 2]));
	bound.up = max(bound.up, diagonal[length - 1] + abs(offDiagonal[length - 2]));

	float norm = max(abs(bound.low), abs(bound.up));
	float error = norm * 2 * length * FLT_EPSILON;
	bound.low -= error;
	bound.up += error;
	return bound;
}

void Lapack::eigens(int n, float* outEigenValues, 
	const float* diag, const float* offDiag, float tolerance,
	int &workSize, float *work, int &intWorkSize, int *intWork) 
{
	int countNeeded = n * 2 * sizeof(Interval) / sizeof(float);
	if (workSize < 0 || intWorkSize < 0) {
		workSize = countNeeded;
		intWorkSize = n;
		return;
	}
	
	if (workSize < countNeeded) throw ArgumentException("Not enough float work space. Need at least 4*N");
	if (intWorkSize < n) throw ArgumentException("Not enough int work space. Need at least N");
		
	Interval* currentIntervals = (Interval*)work;
	Interval* nextIntervals = (Interval*)&work[n * 2];
	
    Interval gerschgorin = getGerschgorinInterval(diag, offDiag, n);
	int numIterations = floor((log(gerschgorin.up - gerschgorin.low) - log(tolerance)) / log(2.0f));
	initIntervals(n, currentIntervals, gerschgorin);
	
	for (int i = 0; i < numIterations; i++) {
		getNumEigensInInterval(intWork, currentIntervals, diag, offDiag, n);		
		recalcEigenIntervals(n, nextIntervals, currentIntervals, intWork, diag, offDiag, tolerance);		
		swap(currentIntervals, nextIntervals);
	}
	
	for (int i = 0; i < n; i++) {
		outEigenValues[i] = (currentIntervals[i].low + currentIntervals[i].up) / 2;
	}
}
